import io
import json
import flask
import torch
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
# from torchvision import transforms as T
from torchvision import transforms, models, datasets
from torch.autograd import Variable

# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False


def load_model():
    """
        加载预先训练好的模型
    """

    # 1.读取官方模型结构
    global model
    # 这里直接加载官方工具包里提供的训练好的模型（代码会自动下载）括号内参数为是否下载模型对应的配置信息
    model = models.resnet18()  # 读取18层模型
    num_ftrs = model.fc.in_features  # 获取模型输出层的特征数参数
    model.fc = nn.Sequential(nn.Linear(num_ftrs, 102))  # 修改模型输出层的类别数，类别数自己根据自己任务来

    # 2.使用我们之前训练好的模型参数进行覆盖
    checkpoint = torch.load('best.pth')  # 读取模型文件
    model.load_state_dict(checkpoint['state_dict'])  # 应用模型参数

    # 3.将模型指定为测试格式
    model.eval()

    # 4.设置是否使用gpu
    if use_gpu:
        model.cuda()


def prepare_image(image, target_size):
    """
        在对任何数据进行预测之前进行图像预处理.
    :param image:       原始图片数据
    :param target_size: 目标图片大小
    :return: 预处理后的图片数据
    """
    # 1.针对不同模型，image的格式可能不同，这里需要统一至RGB格式
    if image.mode != 'RGB':
        image = image.convert("RGB")

    # 2.按照所使用的模型参数，修改输入图片的尺寸，并转为tensor
    image = transforms.Resize(target_size)(image)  # 修改图片的尺寸与我们训练的模型相同
    image = transforms.ToTensor()(image)  # 转为tensor格式

    # 3.对数据进行标准化操作，这里的参数应该和数据集中是对应的，与训练过程中参数保持一致 （RGB三通道）
    image = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)

    # 4.增加一个维度，用于按batch测试   本次这里一次测试一张
    # 训练的时候我们的维度是[batch, channel, H, W]，传入的数据没有batch参数，这里我们手动增加batch=1
    image = image[None]  # 增加batch
    if use_gpu:
        image = image.cuda()
    return Variable(image, volatile=True)  # 不需要求导


# 开启服务   这里的predict只是一个路由名字，可自定义
@app.route("/predict", methods=["POST"])
def predict():
    # 做一个标志，刚开始无图像传入时为false，传入图像时为true
    data = {"success": False}

    # 如果收到请求
    if flask.request.method == 'POST':
        # 判断是否为图像
        if flask.request.files.get("image"):
            # 使用PIL格式读取收到的图像数据
            image = flask.request.files["image"].read()  # 获取数据
            image = Image.open(io.BytesIO(image))  # 二进制数据

            # 利用上面的预处理函数对读入的图像数据进行预处理
            image = prepare_image(image, target_size=(64, 64))

            preds = F.softmax(model(image), dim=1)  # 得到属于各个类别的概率
            results = torch.topk(preds.cpu().data, k=3, dim=1)  # 获取可能性最高的三个类别
            results = (results[0].cpu().numpy(), results[1].cpu().numpy())  # 转到CPU中

            # 字典data新增一个key,value,其中value为list格式，表示预测值
            data['predictions'] = list()

            # 遍历预测结果并将其添加到返回的字典data中
            for prob, label in zip(results[0][0], results[1][0]):
                # label_name = idx2label[str(label)]
                r = {"label": str(label), "probability": float(prob)}
                # 将预测结果添加至data字典
                data['predictions'].append(r)

            # 请求成功标识
            data["success"] = True
    # 将最终结果以json形式返回
    return flask.jsonify(data)


"""
# 返回内容示例
test_json = {
                "status_code": 200,
                "success": {
                            "message": "image uploaded",
                            "code": 200
                        },
                "video":{
                    "video_name":opt['source'].split('/')[-1],
                    "video_path":opt['source'],
                    "description":"1",
                    "length": str(hour)+','+str(minute)+','+str(round(second,4)),
                    "model_object_completed":model_flag
                    }
                    "status_txt": "OK"
                    }
                    response = requests.post(
                        'http://xxx.xxx.xxx.xxx:8090/api/ObjectToKafka/',,
                        data={'json': str(test_json)})
"""

if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    print("Loading model...")
    # 先加载模型
    load_model()
    print("The model has been loaded and is waiting for a request!")
    # 运行API服务，指定服务端口
    app.run(port=5012)
